import asyncio
import ssl
import logging
from concurrent.futures import ThreadPoolExecutor
from fastapi import FastAPI
import uvicorn
from typing import Dict, Optional
from pydantic import BaseModel
import time
import json
import os
import uuid
import threading
from enum import Enum, auto
from collections import deque
import statistics

app = FastAPI()


class ConnectionState(Enum):
    DISCONNECTED = auto()
    CONNECTING = auto()
    CONNECTED = auto()
    RECONNECTING = auto()


class TCPServer:
    def __init__(self, host='0.0.0.0', port=65432, certfile=None, keyfile=None):
        self.host = host
        self.port = port
        self.certfile = certfile
        self.keyfile = keyfile
        self.server = None
        self.clients: Dict[str, asyncio.StreamWriter] = {}
        self.running = False
        self.executor = ThreadPoolExecutor(max_workers=10)
        self.connection_stats = {}
        self.message_stats = {}
        self.start_time = time.time()

    async def handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        client_id = str(uuid.uuid4())
        peername = writer.get_extra_info('peername')
        self.clients[client_id] = writer
        self.connection_stats[client_id] = {
            'connect_time': time.time(),
            'last_activity': time.time(),
            'message_count': 0,
            'bytes_received': 0,
            'bytes_sent': 0
        }

        logging.info(f"Client {peername} connected with ID {client_id}")

        try:
            while self.running:
                try:
                    data = await reader.read(1024)
                    if not data:
                        break

                    self.connection_stats[client_id]['last_activity'] = time.time()
                    self.connection_stats[client_id]['bytes_received'] += len(data)

                    message = data.decode().strip()
                    if message == "HEARTBEAT":
                        writer.write(b"HEARTBEAT_ACK\n")
                        await writer.drain()
                        self.connection_stats[client_id]['bytes_sent'] += 13
                        continue

                    self.connection_stats[client_id]['message_count'] += 1
                    logging.info(f"Received from {peername}: {message}")

                    # Process message and send ACK
                    response = f"ACK: {message}\n"
                    writer.write(response.encode())
                    await writer.drain()
                    self.connection_stats[client_id]['bytes_sent'] += len(response)

                except ConnectionResetError:
                    break
                except Exception as e:
                    logging.error(f"Error handling client {peername}: {e}")
                    break

        finally:
            writer.close()
            await writer.wait_closed()
            del self.clients[client_id]
            logging.info(f"Client {peername} disconnected")

    async def start(self):
        if self.certfile and self.keyfile:
            ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
            ssl_context.load_cert_chain(certfile=self.certfile, keyfile=self.keyfile)
        else:
            ssl_context = None

        self.server = await asyncio.start_server(
            self.handle_client,
            self.host,
            self.port,
            ssl=ssl_context
        )
        self.running = True
        logging.info(f"Server started on {self.host}:{self_port}")

        async with self.server:
            await self.server.serve_forever()

    async def stop(self):
        self.running = False
        if self.server:
            self.server.close()
            await self.server.wait_closed()
        logging.info("Server stopped")

    def get_stats(self):
        uptime = time.time() - self.start_time
        total_messages = sum(stats['message_count'] for stats in self.connection_stats.values())
        total_bytes_received = sum(stats['bytes_received'] for stats in self.connection_stats.values())
        total_bytes_sent = sum(stats['bytes_sent'] for stats in self.connection_stats.values())

        return {
            'uptime': uptime,
            'active_connections': len(self.clients),
            'total_messages': total_messages,
            'total_bytes_received': total_bytes_received,
            'total_bytes_sent': total_bytes_sent,
            'message_rate': total_messages / uptime if uptime > 0 else 0,
            'data_rate': (total_bytes_received + total_bytes_sent) / uptime if uptime > 0 else 0
        }


@app.on_event("startup")
async def startup_event():
    server = TCPServer(
        certfile="server.crt",
        keyfile="server.key"
    )
    app.state.server = server
    asyncio.create_task(server.start())


@app.on_event("shutdown")
async def shutdown_event():
    await app.state.server.stop()


@app.get("/stats")
async def get_stats():
    return app.state.server.get_stats()


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    uvicorn.run(app, host="0.0.0.0", port=8000, ssl_keyfile="server.key", ssl_certfile="server.crt")